package cz.cuni.amis.episodic.dybanem;

import cz.cuni.amis.episodic.bayes.utils.JSmileUtil;

import com.google.common.collect.HashBasedTable;
import com.google.common.collect.Table;

import cz.cuni.amis.episodic.bayes.utils.DynamicNodeProperty;
import cz.cuni.amis.episodic.bayes.utils.NodeSnapshot;
import cz.cuni.amis.episodic.bayes.utils.Utils;

import java.util.*;

import lombok.Data;

import org.apache.commons.lang.ArrayUtils;

import smile.Network;

/**
 * Dynamic Bayes Network class. I haven't found a way how to set "virtual"
 * evidence for temporal nodes in plain JSMILE API. That is why I created this
 * extension. Methods defined here start with getDBN* or setDBN* prefixes to
 * distinguish them from (get|set)Temporal* methods from plain API.
 * 
 * @author ik
 */
@Data
public class DBN {

	Network net;
	int slices;

	public DBN(Network net, int slices) {
		net.setSliceCount(slices);
		this.net = net.unroll().unrolled;
		this.slices = slices;
		// this.net.setBayesianAlgorithm(3); // TODO
	}

	/**
	 * Clears evidence in the network and sets the new evidence at once.
	 * 
	 * @param evidences
	 */
	public void resetEvidence(Collection<? extends Evidence> evidences) {
		getNet().clearAllEvidence();
		for (Evidence e : evidences) {
			e.setEvidence(this);
		}

	}

	public void clearDBNEvidence(String node, int from, int to) {
		for (int i = from; i < to; i++) {
			net.clearEvidence(JSmileUtil.dbnNodeId(node, i));
		}
	}

	/**
	 * Set evidence in time t.
	 * 
	 * @param node
	 * @param t
	 * @param value
	 */
	public void setDBNEvidence(String node, int t, String value) {
		net.setEvidence(JSmileUtil.dbnNodeId(node, t), value);
	}

	public void setDBNEvidence(String node, List<String> hardEvidence) {
		setDBNEvidence(node, hardEvidence, 0);
	}

	public void setDBNEvidence(String node, List<String> hardEvidence,
			int offset) {
		for (int i = 0; i < hardEvidence.size(); i++) {
			setDBNEvidence(node, offset, hardEvidence.get(i));
			offset++;
		}
	}

	public void setDBNVirtualEvidence(String node, int t, double[] evidence) {
		net.setVirtualEvidence(JSmileUtil.dbnNodeId(node, t), evidence);
	}

	public void setDBNVirtualEvidence(String node, List<Double[]> evidence) {
		for (int i = 0; i < evidence.size(); i++) {
			setDBNVirtualEvidence(node, i,
					ArrayUtils.toPrimitive(evidence.get(i)));
		}
	}

	public double[] getDBNNodeValue(String node, int t) {
		return net.getNodeValue(JSmileUtil.dbnNodeId(node, t));
	}

	public List<Double[]> getDBNNodeValues(String node) {
		List<Double[]> vals = new ArrayList<>(slices);
		for (int i = 0; i < slices; i++) {
			vals.add(ArrayUtils.toObject(getDBNNodeValue(node, i)));
		}
		return vals;
	}

	/**
	 * 
	 * @param node
	 * @return first index is time, second outcome
	 */
	public double[][] getDBNNodeValuesArray(String node) {
		return getDBNNodeValuesArray(node, 0, slices);
	}

	/**
	 * 
	 * @param node
	 * @param from
	 *            from time
	 * @param to
	 *            to time
	 * @return first index is time, second outcome
	 */
	public double[][] getDBNNodeValuesArray(String node, int from, int to) {
		double[][] vals = new double[to - from][];
		for (int i = from; i < to; i++) {
			vals[i] = getDBNNodeValue(node, i);
		}
		return vals;
	}

	public double surprise(String evidenceNode, int evidenceSlice,
			String evidence, String targetNode, int targetSlice) {
		// 1) compute distribution on target prior to observation
		net.updateBeliefs();
		double[] prior = getDBNNodeValue(targetNode, targetSlice);

		// 2) compute posterior
		setDBNEvidence(evidenceNode, evidenceSlice, evidence);
		net.updateBeliefs();
		double[] posterior = getDBNNodeValue(targetNode, targetSlice);

		// 3) compute KL divergence between them
		return ProbabilityUtils.kullbackLiebrerDivergence(prior, posterior);
	}

	public NodeSnapshot[] getNodesSnapshot(String[] nodes) {
		NodeSnapshot[] snapshots = new NodeSnapshot[nodes.length];
		for (int i = 0; i < nodes.length; i++) {
			snapshots[i] = getNodeSnapshot(nodes[i]);
		}
		return snapshots;
	}

	public NodeSnapshot getNodeSnapshot(String node) {
		NodeSnapshot snapshot = new NodeSnapshot(node);
		snapshot.setProbability(getDBNNodeValuesArray(node));
		// TODO snapshot.setEntropy(ProbabilityUtils.);
		return snapshot;
	}

	public DynamicNodeProperty[] oneStep(String evidenceNode,
			int evidenceSlice, String evidence, String[] targetNodes,
			int targetSlice) {

		DynamicNodeProperty[] props = new DynamicNodeProperty[targetNodes.length];

		// 1) compute distribution on targets prior to observation
		net.updateBeliefs(); // TODO this can be skipped if properly used in a
								// loop
		for (int i = 0; i < targetNodes.length; i++) {
			DynamicNodeProperty prop = new DynamicNodeProperty(targetSlice,
					targetNodes[i]);
			prop.setFullPrior(getDBNNodeValuesArray(targetNodes[i]));
			props[i] = prop;
		}

		// 2) compute posterior
		setDBNEvidence(evidenceNode, evidenceSlice, evidence);
		net.updateBeliefs();

		for (int i = 0; i < targetNodes.length; i++) {
			DynamicNodeProperty prop = props[i];
			prop.setFullPosterior(getDBNNodeValuesArray(targetNodes[i]));
			// 3) compute KL divergence between prior and posterior, this is
			// surprise
			// 4) compute entropy of posterior
			double[] surprise = new double[slices];
			double[] entropy = new double[slices];
			for (int j = 0; j < entropy.length; j++) {
				entropy[j] = ProbabilityUtils.entropy(getDBNNodeValue(
						targetNodes[i], j));
				surprise[j] = ProbabilityUtils.kullbackLiebrerDivergence(
						prop.getFullPrior()[j], prop.getFullPosterior()[j]);

			}
			prop.setEntropy(entropy);
			prop.setSurprise(surprise);
		}
		return props;
	}

	public Table<Integer, String, Double> getDBNNodeValuesTable(String node) {
		String[] outcomeIds = net.getOutcomeIds(node);
		// List<Integer> rows = new ArrayList<>();
		// for(int i = 0; i < slices; i++) rows.add(i);
		// Table<Integer, String, Double> table = ArrayTable.create(rows,
		// Arrays.asList(outcomeIds));
		Table<Integer, String, Double> table = HashBasedTable.create();

		for (int i = 0; i < slices; i++) {
			double[] vals = getDBNNodeValue(node, i);
			int j = 0;
			for (String id : outcomeIds) {
				table.put(i, id, vals[j]);
				j++;
			}
		}
		return table;
	}

	public int[] getDBNNodeMostProbableValues(String node) {
		double[][] probs = getDBNNodeValuesArray(node);
		int[] outcome = new int[probs.length];
		for (int i = 0; i < probs.length; i++) {
			outcome[i] = Utils.maxIndex(probs[i]);
		}
		return outcome;
	}

	public String[] getDBNNodeMostProbableStringValues(String node) {
		int[] outcomes = getDBNNodeMostProbableValues(node);
		String[] vals = new String[outcomes.length];
		String ids[] = net.getOutcomeIds(node);
		for (int i = 0; i < outcomes.length; i++) {
			vals[i] = ids[outcomes[i]];
		}
		return vals;
	}

	/**
	 * Sets some nodes as targets. This may help in inference.
	 * @param targets
	 */
	public void setDBNTargetNodes(String[] targets) {
		for (String node : targets) {
			for (int i = 0; i < slices; i++) {
				net.setTarget(JSmileUtil.dbnNodeId(node, i), true);
			}
		}
	}
	/*
	 * List<NetSlice> slices = new ArrayList<>();
	 * 
	 * void getSlices() { SliceTemplate template = null; for (int i = 0; i <
	 * net.getSliceCount(); i++) { slices.add(new NetSlice(this, i, template));
	 * } }
	 */
}
